import cv2
import numpy as np

video = r'D:\Documents\data\vtest.avi'
#角点检测参数：
feature_params = dict(maxCorners = 100,
                      qualityLevel = 0.3,
                      minDistance = 7,
                      blockSize = 7)
#L_K光流法参数
lk_params = dict(winSize = (15,15),
                 maxLevel = 2,
                 criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
cap = cv2.VideoCapture(video)
# 随机颜色条
color = np.random.randint(0, 255, (100, 3))
#计算第一特征点
ret,prev = cap.read()
prevGray = cv2.cvtColor(prev,cv2.COLOR_BGR2GRAY)#转成灰度图像
p0 = cv2.goodFeaturesToTrack(prevGray,mask = None, **feature_params)
# 创建一个mask
mask = np.zeros_like(prev)

while True:
    ret,frame = cap.read()
    if not ret:#没读取到当前帧，结束
        break

    gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
    #计算光流
    p1, st, err = cv2.calcOpticalFlowPyrLK(prevGray,gray,p0,None,**lk_params)
    #选取好的跟踪点
    goodPoints = p1[st==1]
    goodPrevPoints = p0[st==1]

    #在结果图像中叠加画出特征点和计算出来的光流向量
    for i, (cur,prev) in enumerate(zip(goodPoints,goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        mask = cv2.line(mask, (x0, y0), (x1, y1), color[i].tolist(), 2)
        frame = cv2.circle(frame, (x0, y0), 5, color[i].tolist(), -1)
    res = cv2.add(frame,mask)

    #显示计算结果图像
    cv2.imshow('result',res)
    key = cv2.waitKey(150)#每一帧间隔30秒
    if key == 27:#按esc退出
        break

    # 更新上一帧
    prevGray = gray.copy()
    p0 = goodPoints.reshape(-1, 1, 2)

cap.release()
cv2.destroyAllWindows()